#include "IKSolver.h"

IKSolver::IKSolver(PhysicsWorld* pw)
{
	this->pw = pw;
}

void IKSolver::updateCCD(PhysicsWorld* pw)
{
	links.clear();
	joints.clear();

	for(map<QString, dJointID>::iterator iter = pw->armJoints.begin(); iter != pw->armJoints.end(); iter++)
	{
		dJointID jointID = iter -> second;
		dJointType jointType = dJointGetType(jointID);

		dVector3 jointPos;
		dJointGetHingeAnchor(jointID, jointPos);

		dVector3 jointAxis;
		dJointGetHingeAxis(jointID, jointAxis);
		dReal jointAngle = dJointGetHingeAngle(jointID);

		Quaternion<double> jointQuat;
		Vector3d axis(jointAxis);
		axis.normalize();
		jointQuat = AngleAxisd(jointAngle, axis);

		Joint joint(jointID, Vector3d(jointPos), jointQuat, Vector3d(jointAxis), jointType);
		joints.push_back(joint);

		dBodyID bodyID = dJointGetBody(jointID, 0);
		const dReal* bodyPos = dBodyGetPosition(bodyID);
		const dReal* bodyQuat = dBodyGetQuaternion(bodyID);

		Link link(bodyID, Vector3d(bodyPos), Quaternion<double>(bodyQuat));
		links.push_back(link);
	}
}

void IKSolver::updateJL(Quaternion<double>& turnedQuat, int i)
{
	//Matrix3d turnedRot = turnedQuat.toRotationMatrix();
	joints.at(i).quat = joints.at(i).quat * turnedQuat;
	Vector3d basePos(joints.at(i).pos);

	for(int j = i; j >= 0; j--)
	{
		Joint joint = joints.at(j);
		Link link = links.at(j);

		Vector3d jointPos(joint.pos);
		jointPos = turnedQuat * (jointPos - basePos) + basePos;
		joints.at(j).pos = jointPos;

		Vector3d linkPos(link.pos);
		linkPos = turnedQuat * (linkPos - basePos) + basePos;
		links.at(j).pos = linkPos;
	}
}

void IKSolver::checkDOFRestrictions(Joint& joint, Quaternion<double>& turnQuat)
{
	dJointID jiontID = joint.id;
	dJointType jointType = joint.type;

	Quaternion<double> jointQuat(joint.quat);
	Quaternion<double> desiredQuat = jointQuat * turnQuat;

	AngleAxisd desiredAA(desiredQuat);
	Vector3d desiredAxis(desiredAA.axis());
	dReal desiredAngle = desiredAA.angle();
	desiredAxis.normalize();

	dVector3 axis;
	dJointGetHingeAxis (jiontID, axis);
	Vector3d jointAxis(axis);

	if((jointAxis - (-desiredAxis)).norm() < POS_THRESH)
		desiredAngle *= -1;

	dReal lo = dJointGetHingeParam(jiontID, dParamLoStop);
	dReal hi = dJointGetHingeParam(jiontID, dParamHiStop);

	desiredAngle = std::max(lo, desiredAngle);
	desiredAngle = std::min(hi, desiredAngle);

	desiredQuat = AngleAxisd(desiredAngle, jointAxis);
	Quaternion<double> jointQuatInv = jointQuat.inverse();
	turnQuat = jointQuatInv * desiredQuat;
}

bool IKSolver::computeCCD(Eigen::Vector3d& endPos)
{
	/// Local Variables ///////////////////////////////////////////////////////////
	Vector3d rootPos,curEnd,desiredEnd,targetVector,curVector,crossResult;
	double cosAngle,turnAngle,turnDeg;
	int	tries  = 0;
	int i = 0;
	Quaternion<double> turnQuat, desiredQuat;

	do
	{
		Joint joint = joints.at(i);
		Link link = links.at(i);

		dJointID jointID = joint.id;
		Vector3d jointPos(joint.pos);
		Quaternion<double> jointQuat(joint.quat);
		dJointType jointType = joint.type;
		Vector3d jointAxis = joint.axis;

		dBodyID linkID = link.id;
		Vector3d linkPos(link.pos);

		rootPos = jointPos;
		curEnd = links.at(0).pos;
		desiredEnd = endPos;

		// SEE IF I AM ALREADY CLOSE ENOUGH
		if ((curEnd - desiredEnd).squaredNorm() > IK_POS_THRESH)
		{
			// CREATE THE VECTOR TO THE CURRENT EFFECTOR POS
			curVector = curEnd - rootPos;
			// CREATE THE DESIRED EFFECTOR POSITION VECTOR
			targetVector = endPos - rootPos;

			
			if(jointType == dJointTypeHinge)
			{
				double dist = targetVector.dot(jointAxis);
				targetVector = endPos - dist * jointAxis - rootPos;
			}

			// NORMALIZE THE VECTORS (EXPENSIVE, REQUIRES A SQRT)
			curVector.normalize();
			targetVector.normalize();

			// THE DOT PRODUCT GIVES ME THE COSINE OF THE DESIRED ANGLE
			cosAngle = targetVector.dot(curVector);

			// IF THE DOT PRODUCT RETURNS 1.0, I DON'T NEED TO ROTATE AS IT IS 0 DEGREES
			if (cosAngle < 0.99999)
			{
				// USE THE CROSS PRODUCT TO CHECK WHICH WAY TO 
				crossResult = curVector.cross(targetVector);
				crossResult.normalize();

				turnAngle = acos(cosAngle);	// GET THE ANGLE
				turnQuat = AngleAxisd(turnAngle, crossResult);

				checkDOFRestrictions(joint, turnQuat);
				updateJL(turnQuat, i);
			}
			if(++i >= joints.size())
				i = 0;
		}

		//qDebug() << curEnd[0] << " " << curEnd[1] << " " << curEnd[2];
		//qDebug() << desiredEnd[0] << " " << desiredEnd[1] << " " << desiredEnd[2] << endl;

	}while((++tries < MAX_IK_TRIES && ((curEnd - desiredEnd).squaredNorm() > IK_POS_THRESH)));

	if (tries == MAX_IK_TRIES)
	{
		return false;
	}
	else
		return true;
}

void IKSolver::reset()
{
	joints.clear();
	links.clear();
}